package goComfyuiApi

import (
	"encoding/json"
	"math"
	"runtime/debug"
	"time"

	"gitee.com/wchpub/fnsvr"
	"github.com/gorilla/websocket"
)

type WsStatusData struct {
	Status PromptResp `json:"status"`
}

type WsExecutingData struct {
	Node        string `json:"node"`
	DisplayNode string `json:"display_node"`
	PromptId    string `json:"prompt_id"`
}

type WsProgressData struct {
	Value    float64 `json:"value"`
	Max      float64 `json:"max"`
	Node     string  `json:"node"`
	PromptId string  `json:"prompt_id"`
}

type WsExecutedData struct {
	Node        string         `json:"node"`
	DisplayNode string         `json:"display_node"`
	PromptId    string         `json:"prompt_id"`
	Output      map[string]any `json:"output"`
}

/*
//{\"type\": \"executed\",
// \"data\": {
// \"output\": {
// \"a_images\": [
// {\"filename\": \"rgthree.compare._temp_axmfm_00005_.png\", \"subfolder\": \"\", \"type\": \"temp\"}
// ],
// \"b_images\": [{\"filename\": \"rgthree.compare._temp_axmfm_00006_.png\", \"subfolder\": \"\", \"type\": \"temp\"}]
}
//
// }}"
*/

type WsExecutionSuccessData struct {
	PromptId  string `json:"prompt_id"`
	Timestamp int64  `json:"timestamp"`
}

func reconnect(url string, clientId string) *websocket.Conn {
	retryInterval := time.Second
	maxInterval := time.Minute
	for {
		conn, _, err := websocket.DefaultDialer.Dial(url+"?clientId="+clientId, nil)
		if err == nil {
			return conn
		}
		time.Sleep(retryInterval)
		fnsvr.Log().Error("reconnect to comfyUI ws failed", "err", err)
		retryInterval = time.Duration(math.Min(float64(retryInterval*2), float64(maxInterval))) //指数退避算法
	}
}

// 处理系统状态类型的消息
func handleStatus(msg []byte) (*WsStatusData, error) {
	var statusMsg struct {
		Data WsStatusData `json:"data"`
	}
	err := json.Unmarshal(msg, &statusMsg)
	if err != nil {
		return nil, err
	}
	return &statusMsg.Data, nil
}

// 执行状态类型的消息
func handleExecuting(msg []byte) (*WsExecutingData, error) {
	var execMsg struct {
		Data WsExecutingData `json:"data"`
	}
	err := json.Unmarshal(msg, &execMsg)
	if err != nil {
		return nil, err
	}
	return &execMsg.Data, nil
}

// 进度类型的消息
func handleProgress(msg []byte) (*WsProgressData, error) {
	var progMsg struct {
		Data WsProgressData `json:"data"`
	}
	err := json.Unmarshal(msg, &progMsg)
	if err != nil {
		return nil, err
	}
	return &progMsg.Data, nil
}

func handleExecuted(msg []byte) (*WsExecutedData, error) {
	var progMsg struct {
		Data WsExecutedData `json:"data"`
	}
	err := json.Unmarshal(msg, &progMsg)
	if err != nil {
		return nil, err
	}
	return &progMsg.Data, nil
}

func handleExecutionSuccess(msg []byte) (*WsExecutionSuccessData, error) {
	var progMsg struct {
		Data WsExecutionSuccessData `json:"data"`
	}
	err := json.Unmarshal(msg, &progMsg)
	if err != nil {
		return nil, err
	}
	return &progMsg.Data, nil
}

type WsCallback interface {
	OnStatus(data *WsStatusData)
	OnExecuting(data *WsExecutingData)
	OnProgress(data *WsProgressData)
	OnExecuted(data *WsExecutedData)
	OnExecutionSuccess(data *WsExecutionSuccessData)
}

type WebsocktListener struct {
	wsConn *websocket.Conn
}

// 监听websocket消息,并阻塞进程
func (m *WebsocktListener) Listen(
	url string, //"ws://10.1.1.126:8188/ws"
	clientId string,
	callback WsCallback,
) {
	defer func() {
		fnsvr.Log().Error("WebsocktListener Listen recovered from panic", "panic", recover(), "stack", debug.Stack())
		if m.wsConn != nil {
			m.wsConn.Close()
		}
	}()
	for {
		m.loop(url, clientId, callback)
	}
}

func (m *WebsocktListener) loop(
	url string,
	clientId string,
	callback WsCallback,
) {

	defer func() {
		if r := recover(); r != nil {
			fnsvr.Log().Error("WebsocktListener.loop recovered from panic", "panic", r, "stack", debug.Stack())
		}
	}()

	if m.wsConn == nil {
		m.wsConn = reconnect(url, clientId)
	}

	_, message, err := m.wsConn.ReadMessage()
	if err != nil {
		fnsvr.Log().Warn("WebsocktListener err close")
		m.wsConn.Close()
		m.wsConn = nil
		return
	}

	var baseMsg struct{ Type string }
	if err := json.Unmarshal(message, &baseMsg); err != nil {
		fnsvr.Log().Error("ApiWebsocktListen unmarshal message error", "err", err)
		return
	}
	switch baseMsg.Type {
	case "progress":
		p, err := handleProgress(message)
		if err != nil {
			fnsvr.Log().Error("ApiWebsocktListen handleProgress error", "err", err)
			return
		}
		callback.OnProgress(p)
	case "executing":
		e, err := handleExecuting(message)
		if err != nil {
			fnsvr.Log().Error("ApiWebsocktListen handleExecuting error", "err", err)
			return
		}
		callback.OnExecuting(e)
	case "status":
		s, err := handleStatus(message)
		if err != nil {
			fnsvr.Log().Error("ApiWebsocktListen handleStatus error", "err", err)
			return
		}
		callback.OnStatus(s)
	case "execution_success":
		e, err := handleExecutionSuccess(message)
		if err != nil {
			fnsvr.Log().Error("ApiWebsocktListen handleExecutionSuccess error", "err", err)
			return
		}
		callback.OnExecutionSuccess(e)

		//fnsvr.Log().Info("WebsocktListener execution_success", "message", string(message))
		//"{\"type\": \"execution_success\", \"data\": {\"prompt_id\": \"e83ba146-c0d8-4d38-bd3b-63176f87dbe0\", \"timestamp\": 1749809269926}}"
	case "crystools.monitor":
		//fnsvr.Log().Info("WebsocktListener crystools.monitor", "message", string(message))
	case "executed":
		e, err := handleExecuted(message)
		if err != nil {
			fnsvr.Log().Error("ApiWebsocktListen handleExecuted error", "err", err)
			return
		}
		callback.OnExecuted(e)

		//fnsvr.Log().Info("WebsocktListener executed", "message", string(message))

	default:
		fnsvr.Log().Warn("WebsocktListener unknown message type", "type", baseMsg.Type)
	}
}

// 创建一个新的WebsocktListener实例
func NewWebsocktListener() *WebsocktListener {
	return &WebsocktListener{}
}
